package com.yuan.localaiagent.config;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Component;
import org.springframework.jdbc.core.JdbcTemplate;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

@Component
public class VectorStoreConfig {


   @Bean("myVectorStore")
   @Primary
   public VectorStore vectorStore(JdbcTemplate jdbcTemplate,
                                  @Qualifier("ollamaEmbeddingModel") EmbeddingModel embeddingModel) {
      PgVectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, embeddingModel)
              .dimensions(768)
              .distanceType(COSINE_DISTANCE)
              .indexType(HNSW)
              .schemaName("ai_vector")
              .vectorTableName("vector_store")
              .maxDocumentBatchSize(10000)
              .build();
      return vectorStore;
   }
}
